from rest_framework import viewsets, status
from rest_framework.decorators import action
from rest_framework.response import Response
from rest_framework.permissions import IsAuthenticated
from django.db.models import Q, Count
from django.utils import timezone
from datetime import datetime, timedelta
import json
import logging

from ..models import SSHSession, SSHCommand, SSHReplayData, AuditRule, AuditAlert
from ..serializers import (
    SSHSessionSerializer, SSHCommandSerializer, SSHReplayDataSerializer,
    AuditRuleSerializer, AuditAlertSerializer
)
from ..jumpserver_audit import jumpserver_audit

logger = logging.getLogger(__name__)


class SSHSessionViewSet(viewsets.ModelViewSet):
    """SSH会话管理"""
    queryset = SSHSession.objects.all()
    serializer_class = SSHSessionSerializer
    permission_classes = [IsAuthenticated]
    lookup_field = 'session_id'  # 使用session_id作为查找字段

    def get_queryset(self):
        queryset = super().get_queryset()

        # 过滤参数
        user_id = self.request.query_params.get('user_id')
        host_id = self.request.query_params.get('host_id')
        status = self.request.query_params.get('status')
        start_date = self.request.query_params.get('start_date')
        end_date = self.request.query_params.get('end_date')

        if user_id:
            queryset = queryset.filter(user_id=user_id)
        if host_id:
            queryset = queryset.filter(host_id=host_id)
        if status:
            queryset = queryset.filter(status=status)
        if start_date:
            queryset = queryset.filter(start_time__gte=start_date)
        if end_date:
            queryset = queryset.filter(start_time__lte=end_date)

        return queryset.order_by('-start_time')

    @action(detail=True, methods=['get'])
    def replay(self, request, session_id=None):
        """获取会话回放数据"""
        try:
            logger.info(f"🔍 回放API被调用，session_id: {session_id}")
            session = self.get_object()
            logger.info(f"🔍 获取到会话: {session.session_id}")

            replay_data = jumpserver_audit.get_session_replay_data(session.session_id)
            logger.info(f"🔍 回放数据获取结果: {replay_data is not None}")

            if replay_data:
                logger.info(f"✅ 回放数据返回成功")
                return Response(replay_data)
            else:
                logger.warning(f"⚠️ 会话回放数据不存在: {session.session_id}")
                return Response(
                    {'error': '会话回放数据不存在'},
                    status=status.HTTP_404_NOT_FOUND
                )
        except Exception as e:
            logger.error(f"❌ 回放API错误: {e}")
            import traceback
            logger.error(f"❌ 错误详情: {traceback.format_exc()}")
            return Response(
                {'error': str(e)},
                status=status.HTTP_500_INTERNAL_SERVER_ERROR
            )

    @action(detail=True, methods=['get'])
    def commands(self, request, pk=None):
        """获取会话命令列表"""
        session = self.get_object()
        commands = session.commands.all().order_by('timestamp')
        serializer = SSHCommandSerializer(commands, many=True)
        return Response(serializer.data)

    @action(detail=False, methods=['get'])
    def statistics(self, request):
        """获取会话统计信息"""
        # 时间范围
        end_date = timezone.now()
        start_date = end_date - timedelta(days=30)

        # 基础统计
        total_sessions = SSHSession.objects.count()
        active_sessions = SSHSession.objects.filter(status='active').count()

        # 最近30天的会话统计
        recent_sessions = SSHSession.objects.filter(
            start_time__gte=start_date
        ).count()

        # 用户活跃度统计
        user_stats = SSHSession.objects.filter(
            start_time__gte=start_date
        ).values('user__username').annotate(
            session_count=Count('id')
        ).order_by('-session_count')[:10]

        # 主机访问统计
        host_stats = SSHSession.objects.filter(
            start_time__gte=start_date
        ).values('host__hostname').annotate(
            session_count=Count('id')
        ).order_by('-session_count')[:10]

        # 风险命令统计
        risk_stats = SSHCommand.objects.filter(
            timestamp__gte=start_date
        ).values('risk_level').annotate(
            count=Count('id')
        )

        return Response({
            'total_sessions': total_sessions,
            'active_sessions': active_sessions,
            'recent_sessions': recent_sessions,
            'user_stats': list(user_stats),
            'host_stats': list(host_stats),
            'risk_stats': list(risk_stats)
        })


class SSHCommandViewSet(viewsets.ModelViewSet):
    """SSH命令管理"""
    queryset = SSHCommand.objects.all()
    serializer_class = SSHCommandSerializer
    permission_classes = [IsAuthenticated]

    def get_queryset(self):
        queryset = super().get_queryset()

        # 过滤参数
        session_id = self.request.query_params.get('session_id')
        user_id = self.request.query_params.get('user_id')
        host_id = self.request.query_params.get('host_id')
        risk_level = self.request.query_params.get('risk_level')
        is_dangerous = self.request.query_params.get('is_dangerous')
        command_search = self.request.query_params.get('command')

        if session_id:
            queryset = queryset.filter(session__session_id=session_id)
        if user_id:
            queryset = queryset.filter(session__user_id=user_id)
        if host_id:
            queryset = queryset.filter(session__host_id=host_id)
        if risk_level:
            queryset = queryset.filter(risk_level=risk_level)
        if is_dangerous:
            queryset = queryset.filter(is_dangerous=is_dangerous.lower() == 'true')
        if command_search:
            queryset = queryset.filter(command__icontains=command_search)

        return queryset.order_by('-timestamp')

    @action(detail=False, methods=['get'])
    def dangerous_commands(self, request):
        """获取危险命令列表"""
        commands = self.get_queryset().filter(is_dangerous=True)
        serializer = self.get_serializer(commands, many=True)
        return Response(serializer.data)


class AuditRuleViewSet(viewsets.ModelViewSet):
    """审计规则管理"""
    queryset = AuditRule.objects.all()
    serializer_class = AuditRuleSerializer
    permission_classes = [IsAuthenticated]

    @action(detail=True, methods=['post'])
    def toggle_active(self, request, pk=None):
        """切换规则启用状态"""
        rule = self.get_object()
        rule.is_active = not rule.is_active
        rule.save()

        return Response({
            'id': rule.id,
            'is_active': rule.is_active,
            'message': f"规则已{'启用' if rule.is_active else '禁用'}"
        })

    @action(detail=False, methods=['post'])
    def test_rule(self, request):
        """测试审计规则"""
        rule_data = request.data
        test_command = rule_data.get('test_command', '')

        if not test_command:
            return Response(
                {'error': '请提供测试命令'},
                status=status.HTTP_400_BAD_REQUEST
            )

        # 创建临时规则对象进行测试
        temp_rule = AuditRule(
            pattern=rule_data.get('pattern', ''),
            rule_type=rule_data.get('rule_type', 'keyword'),
            is_regex=rule_data.get('is_regex', False),
            case_sensitive=rule_data.get('case_sensitive', True)
        )

        # 简单的规则匹配逻辑
        if temp_rule.is_regex:
            import re
            try:
                pattern = re.compile(temp_rule.pattern, 0 if temp_rule.case_sensitive else re.IGNORECASE)
                is_match = bool(pattern.search(test_command))
            except re.error:
                is_match = False
        else:
            if temp_rule.case_sensitive:
                is_match = temp_rule.pattern in test_command
            else:
                is_match = temp_rule.pattern.lower() in test_command.lower()

        return Response({
            'is_match': is_match,
            'test_command': test_command,
            'rule_pattern': temp_rule.pattern
        })


class AuditAlertViewSet(viewsets.ModelViewSet):
    """审计告警管理"""
    queryset = AuditAlert.objects.all()
    serializer_class = AuditAlertSerializer
    permission_classes = [IsAuthenticated]

    def get_queryset(self):
        queryset = super().get_queryset()

        # 过滤参数
        status_filter = self.request.query_params.get('status')
        risk_level = self.request.query_params.get('risk_level')
        user_id = self.request.query_params.get('user_id')

        if status_filter:
            queryset = queryset.filter(status=status_filter)
        if risk_level:
            queryset = queryset.filter(risk_level=risk_level)
        if user_id:
            queryset = queryset.filter(session__user_id=user_id)

        return queryset.order_by('-created_at')

    @action(detail=True, methods=['post'])
    def handle_alert(self, request, pk=None):
        """处理告警"""
        alert = self.get_object()
        action = request.data.get('action')  # confirm, ignore, resolve
        note = request.data.get('note', '')

        if action not in ['confirmed', 'ignored', 'resolved']:
            return Response(
                {'error': '无效的处理动作'},
                status=status.HTTP_400_BAD_REQUEST
            )

        alert.status = action
        alert.handler = request.user
        alert.handle_time = timezone.now()
        alert.handle_note = note
        alert.save()

        return Response({
            'id': alert.id,
            'status': alert.status,
            'message': f"告警已{action}"
        })

    @action(detail=False, methods=['get'])
    def pending_count(self, request):
        """获取待处理告警数量"""
        count = self.get_queryset().filter(status='pending').count()
        return Response({'pending_count': count})


class SSHReplayViewSet(viewsets.ReadOnlyModelViewSet):
    """SSH回放数据管理"""
    queryset = SSHReplayData.objects.all()
    serializer_class = SSHReplayDataSerializer
    permission_classes = [IsAuthenticated]

    def get_queryset(self):
        queryset = super().get_queryset()
        session_id = self.request.query_params.get('session_id')

        if session_id:
            queryset = queryset.filter(session__session_id=session_id)

        return queryset.order_by('timestamp')

    @action(detail=False, methods=['get'])
    def export_session(self, request):
        """导出会话回放数据"""
        session_id = request.query_params.get('session_id')

        if not session_id:
            return Response(
                {'error': '请提供会话ID'},
                status=status.HTTP_400_BAD_REQUEST
            )

        try:
            replay_data = jumpserver_audit.get_session_replay_data(session_id)

            # 生成导出文件名
            session = SSHSession.objects.get(session_id=session_id)
            filename = f"ssh_replay_{session.user.username}_{session.host.hostname}_{session.start_time.strftime('%Y%m%d_%H%M%S')}.json"

            response = Response(replay_data)
            response['Content-Disposition'] = f'attachment; filename="{filename}"'
            response['Content-Type'] = 'application/json'

            return response

        except SSHSession.DoesNotExist:
            return Response(
                {'error': '会话不存在'},
                status=status.HTTP_404_NOT_FOUND
            )
        except Exception as e:
            return Response(
                {'error': str(e)},
                status=status.HTTP_500_INTERNAL_SERVER_ERROR
            )
